package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/pelletier/go-toml"
)

type Data struct {
	Dirs     []string
	Packages []string
	Branchs  []string
}

func main() {
	c := flag.String("c", "", "toml content")
	flag.Parse()
	data := &Data{}
	if err := toml.Unmarshal([]byte(*c), data); err != nil {
		logErr(err, "toml Unmarshal")
		return
	}
	if len(data.Packages) != len(data.Branchs) {
		log.Println("Packages and Branchs len is not equal")
		return
	}
	packMap := make(map[string]string)
	for i := 0; i < len(data.Packages); i++ {
		packMap[data.Packages[i]] = data.Branchs[i]
	}
	if len(packMap) == 0 {
		log.Println("packMap is empty")
		return
	}
	for _, dir := range data.Dirs {
		content, err := readGoMod(dir, packMap)
		if err != nil {
			logErr(err, "readGoMod")
			return
		}
		if err := writeGoMod(dir, content); err != nil {
			logErr(err, "writeGoMod")
			continue
		}
		gomodTidy(dir)
		log.Printf("update %s ok", dir)
	}
}

func logErr(err error, msg string) {
	if err == nil {
		return
	}
	log.Printf("[ERR] - %s err: %v", msg, err)
}

func readGoMod(dir string, packMap map[string]string) (content []byte, err error) {
	file, err := os.Open(dir + "/go.mod")
	if err != nil {
		logErr(err, "readGoMod Open")
		return
	}
	defer file.Close()
	reader := bufio.NewReader(file)
	content = make([]byte, 0)
	var (
		line    []byte
		readErr error
	)
	for readErr == nil {
		if line, _, readErr = reader.ReadLine(); readErr != nil {
			if readErr != io.EOF {
				err = readErr
				return
			}
			break
		}
		l := strings.TrimSpace(string(line))
		idx := strings.LastIndex(l, " ")
		if idx > 0 {
			pack := l[:idx]
			if branch, ok := packMap[pack]; ok {
				line = []byte(pack + " " + branch)
			}
		}
		line = append(line, '\n')
		content = append(content, line...)
	}
	return
}

func writeGoMod(dir string, content []byte) (err error) {
	file, err := os.OpenFile(dir+"/go.mod", os.O_WRONLY|os.O_TRUNC, os.ModePerm)
	if err != nil {
		logErr(err, "writeGoMod OpenFile")
		return
	}
	defer file.Close()
	writer := bufio.NewWriter(file)
	if _, err = writer.Write(content); err != nil {
		logErr(err, "writeGoMod Write")
		return
	}
	err = writer.Flush()
	return
}

func gomodTidy(dir string) {
	exec.Command("sh", "-c", fmt.Sprintf("cd %s; go mod tidy; cd ..", dir)).Run()
}
